// Copyright (c) Lawrence Livermore National Security, LLC and other VisIt
// Project developers.  See the top-level LICENSE file for dates and other
// details.  No copyright assignment is required to contribute to VisIt.

// ************************************************************************* //
//  File: avtExtractPointFunction2DFilter.C
// ************************************************************************* //

#include <avtExtractPointFunction2DFilter.h>

#include <float.h>

#ifdef PARALLEL
#include <avtParallel.h>
#include <mpi.h>
#endif

#include <vtkCellData.h>
#include <vtkDataArray.h>
#include <vtkDoubleArray.h>
#include <vtkFieldData.h>
#include <vtkIntArray.h>
#include <vtkRectilinearGrid.h>

#include <vtkDataSetWriter.h>

// ****************************************************************************
//  Method: avtExtractPointFunction2DFilter constructor
//
//  Programmer: ghweber -- generated by xml2avt
//  Creation:   Tue Jun 4 10:46:56 PDT 2013
//
// ****************************************************************************

avtExtractPointFunction2DFilter::avtExtractPointFunction2DFilter()
{
    range[0] = 0;
    range[1] = 0;
    spatialExtents[0] = 0;
    spatialExtents[1] = 0;
    spatialExtents[2] = 0;
    spatialExtents[3] = 0;
    spatialExtents[4] = 0;
    spatialExtents[5] = 0;
}


// ****************************************************************************
//  Method: avtExtractPointFunction2DFilter destructor
//
//  Programmer: ghweber -- generated by xml2avt
//  Creation:   Tue Jun 4 10:46:56 PDT 2013
//
//  Modifications:
//
// ****************************************************************************

avtExtractPointFunction2DFilter::~avtExtractPointFunction2DFilter()
{
}


// ****************************************************************************
//  Method:  avtExtractPointFunction2DFilter::Create
//
//  Programmer: ghweber -- generated by xml2avt
//  Creation:   Tue Jun 4 10:46:56 PDT 2013
//
// ****************************************************************************

avtFilter *
avtExtractPointFunction2DFilter::Create()
{
    return new avtExtractPointFunction2DFilter();
}


// ****************************************************************************
//  Method:      avtExtractPointFunction2DFilter::SetAtts
//
//  Purpose:
//      Sets the state of the filter based on the attribute object.
//
//  Arguments:
//      a        The attributes to use.
//
//  Programmer: ghweber -- generated by xml2avt
//  Creation:   Tue Jun 4 10:46:56 PDT 2013
//
// ****************************************************************************

void
avtExtractPointFunction2DFilter::SetAtts(const AttributeGroup *a)
{
    atts = *(const ExtractPointFunction2DAttributes*)a;
}


// ****************************************************************************
//  Method: avtExtractPointFunction2DFilter::Equivalent
//
//  Purpose:
//      Returns true if creating a new avtExtractPointFunction2DFilter with the given
//      parameters would result in an equivalent avtExtractPointFunction2DFilter.
//
//  Programmer: ghweber -- generated by xml2avt
//  Creation:   Tue Jun 4 10:46:56 PDT 2013
//
// ****************************************************************************

bool
avtExtractPointFunction2DFilter::Equivalent(const AttributeGroup *a)
{
    return (atts == *(ExtractPointFunction2DAttributes*)a);
}


// ****************************************************************************
//  Method: avtExtractPointFunction2DFilter::PreExecute
//
//  Purpose:
//      Called before "Execute". Initialize variable we use to keep track of
//      value range.
//
//  Programmer: Gunther H. Weber 
//  Creation:   June 6, 2013
//
//  Modifications:
//
// ****************************************************************************

void
avtExtractPointFunction2DFilter::PreExecute(void)
{
    avtPluginDataTreeIterator::PreExecute();
    range[0] =  FLT_MAX;
    range[1] = -FLT_MAX;
    spatialExtents[0] = 0;
    spatialExtents[1] = 0;
    spatialExtents[2] = 0;
    spatialExtents[3] = 0;
    spatialExtents[4] = 0;
    spatialExtents[5] = 0;
}

// ****************************************************************************
//  Method: avtExtractPointFunction2DFilter::Execute
//
//  Purpose:
//      Extracts the point function in a localtion (or a list of locations)
//
//  Programmer: ghweber -- generated by xml2avt
//  Creation:   Tue Jun 4 10:46:56 PDT 2013
//
//  Modifications:
//    Kathleen Biagas, Tue Jul  8 14:52:03 PDT 2014
//    Moved avtRealDims code into FIXME comment, to suppress warnings until
//    such a time as they are used.
//
//    Eric Brugger, Thu Jul 24 13:27:16 PDT 2014
//    Modified the class to work with avtDataRepresentation.
//
// ****************************************************************************

void
avtExtractPointFunction2DFilter::Execute()
{
    // Check consistency
    if (atts.GetI().size() != atts.GetJ().size())
    {
        EXCEPTION1(ImproperUseException,
                "I and J arrays need to have the same number of elements.");
    }

    const size_t num_functions = atts.GetI().size();
    std::vector<double*> function_data(num_functions, 0);
    bool klInformationSet = false;
    double dvpar = 1.0;
    double dmu = 1.0;
    int v_base_index[2] = { 0, 0 };
    int v_dims[2] = { 0, 0 };

    // Get the input data tree
    avtDataTree_p in_tree = GetInputDataTree();

    if (in_tree->IsEmpty())
    {
        SetOutputDataTree(in_tree);
        return;
    }

    int nLeaves = 0;
    vtkDataSet **leaves = in_tree->GetAllLeaves(nLeaves);

    // Get all data sets
    // FIXME: Structured grids should be simple as well
    vtkRectilinearGrid *rgrid = 0; // HACK: This just uses the last instance
    vtkDataArray *data = 0; // HACK: This just uses the last instance
    for (int i = 0; i < nLeaves; ++i)
    {
        // Ensure that we are working on rectilinear grid
        rgrid = dynamic_cast<vtkRectilinearGrid*>(leaves[i]);
        if (!rgrid)
            EXCEPTION1(ImproperUseException,
                       "Can only extract point function for a rectilinear grid.");

        // Dimensions of data set
        int dims[3];
        rgrid->GetDimensions(dims);
        // We want number of cells as grid dimension, not number of samples
        for (int d=0; d<3; ++d)
            dims[d]--;

        vtkIntArray *arr = dynamic_cast<vtkIntArray*>(rgrid->GetFieldData()->GetArray("base_index"));
        int base_index[3] = { 0, 0, 0 };
        if (arr)
            for (int d = 0; d < 3; ++d)
                base_index[d] = arr->GetValue(d);

        int i_min = base_index[0];
        int i_max = base_index[0] + dims[0] - 1;
        int j_min = base_index[1];
        int j_max = base_index[1] + dims[1] - 1;

        // FIXME: Take avtRealDims into account
        // arr = dynamic_cast<vtkIntArray*>(rgrid->GetFieldData()->GetArray("avtRealDims"));
        // int avtRealDims[6] = { 0, 0, 0, 0, 0, 0 };
        // if (arr)
        //     for (int d = 0; d < 6; ++d)
        //         avtRealDims[d] = arr->GetValue(d);

        vtkDoubleArray *dx_arr = dynamic_cast<vtkDoubleArray*>(rgrid->GetFieldData()->GetArray("dx_array"));
        vtkIntArray *v_base_index_arr = dynamic_cast<vtkIntArray*>(rgrid->GetFieldData()->GetArray("v_base_index"));
        vtkIntArray *v_dims_arr = dynamic_cast<vtkIntArray*>(rgrid->GetFieldData()->GetArray("v_dims"));

        const char *justTheVar = pipelineVariable + strlen("operators/ExtractPointFunction2D/");
 
        if (!klInformationSet)
        {
           if (dx_arr)
            {
                dvpar = dx_arr->GetValue(0);
                dmu = dx_arr->GetValue(1);
            }
            if (v_base_index_arr && v_dims_arr)
            {
                for (int i = 0; i < 2; ++i)
                {
                    v_base_index[i] = v_base_index_arr->GetValue(i);
                    v_dims[i] = v_dims_arr->GetValue(i);
                }
            }
            else
            {
                EXCEPTION1(ImproperUseException,
                        "Internal error: Velocity base index and dimensions not set by database plugin.");
            }
        }
        else
        {
           if (dx_arr)
            {
                if (dvpar != dx_arr->GetValue(0))
                    EXCEPTION1(ImproperUseException, "Internal error: dvpar mismatch.");
                if (dmu != dx_arr->GetValue(1))
                    EXCEPTION1(ImproperUseException, "Internal error: dmu mismatch.");
            }
            if (v_base_index_arr && v_dims_arr)
            {
                for (int i = 0; i < 2; ++i)
                {
                    if (v_base_index[i] != v_base_index_arr->GetValue(i))
                        EXCEPTION1(ImproperUseException, "Internal error: v_base_index mismatch.");
                    if (v_dims[i] != v_dims_arr->GetValue(i))
                        EXCEPTION1(ImproperUseException, "Internal error: v_dims mismatch.");
                }
            }
        }

        data = rgrid->GetCellData()->GetArray(justTheVar); // FIXME HACK: Outside for loop to ensure that processor 0 has a data pointer

        if (!data)
        {
            EXCEPTION1(VisItException, "Internal error: Could not get data for array variable (maybe due to an operator requesting ghost zones).");
        }

        const std::vector<int> &i_vals = atts.GetI();
        const std::vector<int> &j_vals = atts.GetJ();
        for (size_t curr_tuple = 0; curr_tuple < i_vals.size(); ++curr_tuple)
            if (i_min <= i_vals[curr_tuple] && i_vals[curr_tuple] <= i_max && j_min <= j_vals[curr_tuple] && j_vals[curr_tuple] <= j_max)
            {
                int ijk_f[3] = { i_vals[curr_tuple] - base_index[0], j_vals[curr_tuple] - base_index[1], 0 };
                vtkIdType id_f = rgrid->ComputeCellId(ijk_f);
                function_data[curr_tuple] = new double[v_dims[0]*v_dims[1]];
                for (vtkIdType comp = 0; comp < v_dims[0]*v_dims[1]; ++comp)
                    function_data[curr_tuple][comp] = data->GetComponent(id_f, comp);
            }
    }

#ifdef PARALLEL
    // Determine (on processor 0) where data resides
    std::vector<int> function2proc_map(function_data.size(), -1);
    for (size_t curr_func = 0; curr_func < function_data.size(); ++curr_func)
        if (function_data[curr_func]) function2proc_map[curr_func] = PAR_Rank();
    std::vector<int> tmp_function2proc_map(function2proc_map);
    MPI_Reduce(&tmp_function2proc_map[0], &function2proc_map[0], function2proc_map.size(), MPI_INT, MPI_MAX, 0, VISIT_MPI_COMM);

    // HACK: Send all data to processor 0
    // FIXME: Better distribution among processors
    // FIXME: Non-blocking send/receive
    if (PAR_Rank() == 0)
    {
        for (size_t curr_func = 0; curr_func < function_data.size(); ++curr_func)
            if (function2proc_map[curr_func] >= 1)
            {
                function_data[curr_func] = new double[v_dims[0]*v_dims[1]];
                MPI_Status status;
                MPI_Recv(function_data[curr_func], v_dims[0]*v_dims[1], MPI_DOUBLE, function2proc_map[curr_func], MPI_ANY_TAG, VISIT_MPI_COMM, &status);
            }
    }
    else
    {
        for (size_t curr_func = 0; curr_func < function_data.size(); ++curr_func)
            if (function_data[curr_func])
                MPI_Send(function_data[curr_func], v_dims[0]*v_dims[1], MPI_DOUBLE, 0, 0, VISIT_MPI_COMM);
    }

    if (PAR_Rank() == 0)
    {
#endif
        vtkRectilinearGrid *ogrid = vtkRectilinearGrid::New();
        ogrid->SetDimensions(v_dims[0]+1, v_dims[1]+1, function_data.size());
        vtkDataArray *xCoords = rgrid->GetXCoordinates()->NewInstance();
        xCoords->SetNumberOfTuples(v_dims[0]+1);
        for (int k = 0; k < v_dims[0]+1; ++k)
            xCoords->SetTuple1(k, (v_base_index[0]+k)*dvpar);
        ogrid->SetXCoordinates(xCoords);
        xCoords->Delete();
        vtkDataArray *yCoords = rgrid->GetXCoordinates()->NewInstance();
        yCoords->SetNumberOfTuples(v_dims[1]+1);
        for (int l = 0; l < v_dims[1]+1; ++l)
            yCoords->SetTuple1(l, (v_base_index[1]+l)*dmu);
        ogrid->SetYCoordinates(yCoords);
        yCoords->Delete();
        vtkDataArray *zCoords = rgrid->GetXCoordinates()->NewInstance();
        zCoords->SetNumberOfTuples(function_data.size());
        int fd_size = static_cast<int>(function_data.size());
        for (int f = 0; f < fd_size; ++f)
            zCoords->SetTuple1(f, f);
        ogrid->SetZCoordinates(zCoords);
        zCoords->Delete();

        // FIXME: v_base_index, dvpar and dmu are only defined if there is any data on rank 0
        spatialExtents[0] = v_base_index[0] * dvpar;
        spatialExtents[1] = (v_base_index[0] + v_dims[0]) * dvpar;
        spatialExtents[2] = v_base_index[1] * dmu;
        spatialExtents[3] = (v_base_index[1] + v_dims[1]) * dmu;
        spatialExtents[4] = 0;
        spatialExtents[5] = function_data.size() - 1;

        // FIXME: data is only defined if there is any data on rank 0
        vtkDataArray *odata = data->NewInstance();
        odata->SetNumberOfComponents(1);
        odata->SetNumberOfTuples(v_dims[0]*v_dims[1]*function_data.size());
        for (int k = 0; k < v_dims[0]; ++k)
            for (int l = 0; l < v_dims[1]; ++l)
                for (int f = 0; f < fd_size; ++f)
                {
                    int ijk_t[3] = { k, l, f };
                    vtkIdType id_t = ogrid->ComputeCellId(ijk_t);
                    double val = function_data[f] ? function_data[f][k*v_dims[1]+l] : 0;
                    odata->SetTuple1(id_t, val);
                    if (val < range[0]) range[0] = val;
                    if (val > range[1]) range[1] = val;
                }
        odata->SetName(outVarName.c_str());
        odata->SetName(outVarName.c_str());
        ogrid->GetCellData()->SetScalars(odata);
        odata->Delete();

#if 0
        vtkDataSetWriter *wrtr = vtkDataSetWriter::New();
        wrtr->SetFileTypeToASCII();
        wrtr->SetInputData(ogrid);
        wrtr->SetFileName("debug.vtk");
        wrtr->Write();
#endif

        SetOutputDataTree(new avtDataTree(ogrid, 1));
#ifdef PARALLEL
    }
    else
    {
        SetOutputDataTree(new avtDataTree());
    }
#endif

    for (std::vector<double*>::iterator it = function_data.begin(); it != function_data.end(); ++it)
        delete[] *it;
}


// ****************************************************************************
//  Method: avtExtractPointFunction2DFilter::PostExecute
//
//  Purpose:
//      Called after "Execute". Update value range.
//
//  Programmer: Gunther H. Weber 
//  Creation:   June 6, 2013
//
//  Modifications:
//
// ****************************************************************************

void
avtExtractPointFunction2DFilter::PostExecute(void)
{
    avtPluginDataTreeIterator::PostExecute();
    if (outVarName != "")
    {
        avtDataAttributes &atts = GetOutput()->GetInfo().GetAttributes();
        atts.GetThisProcsOriginalDataExtents(outVarName.c_str())->Set(range);
        atts.GetThisProcsActualDataExtents(outVarName.c_str())->Set(range);
        atts.GetOriginalSpatialExtents()->Clear();
        atts.GetThisProcsOriginalSpatialExtents()->Set(spatialExtents);
    }
}


// ****************************************************************************
//  Method: avtExtractPointFunction2DFilter::ModifyContract
//
//  Purpose:
//      Creates a contract the removes the operator-created-expression.
//
//  Programmer: ghweber -- generated by xml2avt
//  Creation:   Tue Jun 4 10:46:56 PDT 2013
//
// ****************************************************************************

avtContract_p
avtExtractPointFunction2DFilter::ModifyContract(avtContract_p in_contract)
{
    avtContract_p rv;
    if (strncmp(pipelineVariable, "operators/ExtractPointFunction2D/", strlen("operators/ExtractPointFunction2D/")) == 0)
    {
        outVarName = pipelineVariable;
        const char *justTheVar = pipelineVariable + strlen("operators/ExtractPointFunction2D/");
        avtDataRequest_p dr = new avtDataRequest(in_contract->GetDataRequest(), justTheVar);
        rv = new avtContract(in_contract, dr);
    }
    else
    {
        rv = new avtContract(in_contract);
    }
    return rv;
}


// ****************************************************************************
//  Method: avtExtractPointFunction2DFilter::UpdateDataObjectInfo
//
//  Purpose:
//      Tells output that we have a new variable.
//
//  Programmer: ghweber -- generated by xml2avt
//  Creation:   Tue Jun 4 10:46:56 PDT 2013
//
// ****************************************************************************

void
avtExtractPointFunction2DFilter::UpdateDataObjectInfo(void)
{
    avtDataAttributes &data_atts = GetOutput()->GetInfo().GetAttributes();
    data_atts.SetSpatialDimension(atts.GetI().size() == 1 ? 2 : 3);
    data_atts.SetTopologicalDimension(atts.GetI().size() == 1 ? 2 : 3);
    data_atts.AddFilterMetaData("ExtractPointFunction2D");

    GetOutput()->GetInfo().GetValidity().SetPointsWereTransformed(true);
    GetOutput()->GetInfo().GetAttributes().SetCanUseTransform(false);
    GetOutput()->GetInfo().GetAttributes().SetCanUseInvTransform(false);
    GetOutput()->GetInfo().GetValidity().InvalidateZones();
    GetOutput()->GetInfo().GetValidity().InvalidateSpatialMetaData();
    GetOutput()->GetInfo().GetValidity().InvalidateDataMetaData();

   if (outVarName != "")
   {
       if (! data_atts.ValidVariable(outVarName))
       {
           data_atts.AddVariable(outVarName.c_str());
           data_atts.SetActiveVariable(outVarName.c_str());
           data_atts.SetVariableDimension(1);
           data_atts.SetVariableType(AVT_SCALAR_VAR);
           data_atts.SetCentering(AVT_ZONECENT);
       }
   }
}
